from gensim import utils
from gensim.models.doc2vec import TaggedDocument
from random import shuffle

from config import data_path


class TaggedSentence(object):
    def __init__(self, sources):
        self.sources = sources
        self.sentences = []

    def __iter__(self):
        for prefix, source in self.sources.items():
            with utils.smart_open(source) as fin:
                for item_no, line in enumerate(fin):
                    yield TaggedDocument(utils.to_unicode(line).split(), [prefix + '_%s' % item_no])

    def to_list(self):
        self.sentences = []
        for prefix, source, in self.sources.items():
            with utils.smart_open(source) as fin:
                for item_no, line in enumerate(fin):
                    self.sentences.append(TaggedDocument(utils.to_unicode(line).split(), [prefix + '_%s' % item_no]))
        return self.sentences

    def sentences_perm(self):
        shuffle(self.sentences)
        return self.sentences

if __name__ == '__main__':
    pass
    sources = {
        '1': data_path('1.txt'),
        '2': data_path('2.txt')
    }
    sentences = TaggedSentence(sources)
    for sentence, prefix in sentences:
        print(sentence, prefix)
